using System;
using System.Reflection;
using NBody;
using NBody.InitialConditions;
using NBody.Cosmology;
using OptionParser;
using Gsl;
using Fftw;
using DataFile;

[assembly: AssemblyTitle ("NBodyCube2")]
[assembly: AssemblyVersion ("1.0.*")]
[assembly: AssemblyDescription ("Generate comoving initial conditions")]
[assembly: AssemblyCopyright ("2006 Joseph D. MacMillan")]

public class InitOptions : OutputGetOptions
{
    [Option ("-np", "Number of particles per side of initial cube")]
    public int Np = 32;
    
    [Option ("-n", "Number of displacements")]
    public int N = 32;
    
    [Option ("-l", "Comoving length of initial cube (Mpc)")]
    public double L = 10.0;
    
    [Option ("-z", "Initial redshift")]
    public double Z = 30.0;
    
    [Option ("-seed", "Random seed")]
    public int Seed = 0;
}

public class NBodyInit
{
    static double A;
    static double h = 0.5;
    static double H0 = 2.4108e-3; // for kpc, 10^8 M_sun, etc, with h = 0.5
    static double sigma_8 = 0.63;
    static double R0 = 8.0 / h;
    static double Z;
    
    public static void Main(string[] args)
    {
        InitOptions opts = new InitOptions();
        opts.ProcessArgs(args);
        
        // set up units in comoving Mpc
        Z = opts.Z;
        double Length = opts.L;
        
        // normalize power spectrum
        Gsl.Integration integ = new Gsl.Integration(1000);
        integ.Function = Sigma0;        
        integ.ExecuteQAGIU(0.0);
        double sigma2 = 1.0 / 2.0 / Math.PI / Math.PI * integ.Result;
        A = sigma_8 * sigma_8 / sigma2;
        Console.Error.WriteLine("Power spectrum amplitude = {0}", A);
        // scale power spectrum normalization to redshift
        A *= 1.0 / (1.0 + Z) / (1.0 + Z);
        Console.Error.WriteLine("  Scaled amplitude = {0}", A);
        
        Console.Error.WriteLine("Generating cube ... ");
        RandomGaussian random;
        if (opts.Seed == 0)
            random = new RandomGaussian();
        else
            random = new RandomGaussian(opts.Seed);
        double delta_k = 2.0 * Math.PI / Length;
            
        Fftw.Complex[] fx = new Fftw.Complex[opts.N * opts.N * (opts.N/2 + 1)];
        Fftw.Complex[] fy = new Fftw.Complex[opts.N * opts.N * (opts.N/2 + 1)];
        Fftw.Complex[] fz = new Fftw.Complex[opts.N * opts.N * (opts.N/2 + 1)];
        
        double a, b;
        double kx, ky, kz;
        int c = 0;
        for (int i = 0; i < opts.N; i++)
        {
            kx = delta_k * i;
            for (int j = 0; j < opts.N; j++)
            {
                ky = delta_k * j;
                for (int k = 0; k < opts.N/2 + 1; k++)
                {
                    kz = delta_k * k;
                       
                    double kk = Math.Sqrt(kx*kx + ky*ky + kz*kz);
                        
                    double r1 = random.NextDouble();
                    double r2 = random.NextDouble();
                        
                    if (kk == 0.0)
                    {
                        a = 0.0;
                        b = 0.0;
                    }
                    else
                    {
                        double p = P(kk);
                        a = Math.Sqrt(p) / kk / kk * r1;
                        b = Math.Sqrt(p) / kk / kk * r2;
                    }
                        
                    fx[c].Re = 0.5 * kx * b;
                    fx[c].Im = 0.5 * kx * a;
                      
                    fy[c].Re = 0.5 * ky * b;
                    fy[c].Im = 0.5 * ky * a;
                      
                    fz[c].Re = 0.5 * kz * b;
                    fz[c].Im = 0.5 * kz * a;
                        
                    c++;
                }
            }
        }    
            
        PlanComplexToReal3D planx = new PlanComplexToReal3D(opts.N, opts.N, opts.N, fx, Plan.FFTW_ESTIMATE);
        planx.Execute();
        double[] sx = planx.RealOutput;
            
        PlanComplexToReal3D plany = new PlanComplexToReal3D(opts.N, opts.N, opts.N, fy, Plan.FFTW_ESTIMATE);
        plany.Execute();
        double[] sy = plany.RealOutput;
            
        PlanComplexToReal3D planz = new PlanComplexToReal3D(opts.N, opts.N, opts.N, fz, Plan.FFTW_ESTIMATE);
        planz.Execute();
        double[] sz = planz.RealOutput;
        
        while (opts.Np != opts.N)
        {
            Resample(ref sx, ref sy, ref sz, opts.N);
            Console.Error.WriteLine("New sx.Length = {0}", sx.Length);
            opts.N /= 2;
        }
        
        // Generate a few simple stats about the displacements
        int ntotal = opts.N * opts.N * opts.N;
            
        double avg = 0.0;
        double max = 0.0;
        for (int i = 0; i < ntotal; i++)
        {
            double dd = Math.Sqrt(sx[i]*sx[i] + sy[i]*sy[i] + sz[i]*sz[i]);
            avg += dd;
            if (dd > max) max = dd;
        }
        avg /= (double)ntotal;
        double grid = Length / (double)opts.N;
            
        Console.Error.WriteLine("Grid spacing: {0:f6}", grid);
        Console.Error.WriteLine("Avg displacements: {0:f6}", avg / (1.0 + Z));
        Console.Error.WriteLine("Max displacements: {0:f6}", max / (1.0 + Z));
        Console.Error.WriteLine("Redhsift should be around z = {0}", 1.0 / Math.Sqrt(grid / (1.0 + Z) / max) - 1.0);
                       
        // now apply displacements to particles on a cubic grid
        double time = 2.0 / 3.0 / H0 * Math.Pow(1.0 + Z, -3.0 / 2.0);
        NBodySystem s = new NBodySystem(ntotal, time);
        double delta_x = Length / (double)opts.N;
        double Mtot = 3.0 * H0 * H0 / 8.0 / Math.PI * Length * Length * Length * 1e9; // convert length units to kpc        
        double m = Mtot / (double)ntotal;
        c = 0;
        for (int i = 0; i < opts.N; i++)
        {
            double x = delta_x * i;
            for (int j = 0; j < opts.N; j++)
            {
                double y = delta_x * j;
                for (int k = 0; k < opts.N; k++)
                {
                    double z = delta_x * k;
                    
                    double xs, ys, zs;
                    double vx, vy, vz;
                    
                    xs = x + 1.0 / (1.0 + Z) * sx[c];
                    while (xs > Length) xs -= Length; while (xs < 0.0) xs += Length;
                    ys = y + 1.0 / (1.0 + Z) * sy[c];
                    while (ys > Length) ys -= Length; while (ys < 0.0) ys += Length;
                    zs = z + 1.0 / (1.0 + Z) * sz[c];
                    while (zs > Length) zs -= Length; while (zs < 0.0) zs += Length;
                    
                    vx = H0 / Math.Sqrt(1.0 + Z) * sx[c];
                    vy = H0 / Math.Sqrt(1.0 + Z) * sy[c];
                    vz = H0 / Math.Sqrt(1.0 + Z) * sz[c];
                      
                    s.AddParticle(new Particle(m, xs, ys, zs, vx, vy, vz));
                    
                    c++;
                }
            }
        }
        
        for (int i = 0; i < s.NumParts; i++)
        {
            // now scale positions and velocites to correct units
            double sc = 1000.0;
            s[i].SetPosition(sc * s[i].X, sc * s[i].Y, sc * s[i].Z);
            s[i].SetVelocity(sc * s[i].Vx, sc * s[i].Vy, sc * s[i].Vz);
        }
        
        s.Write(opts.OutFile);
    }
    
    static double Sigma0(double k, object o)
    {
        double q = k;
        double x = k * R0;
        double w = 3.0 * (Math.Sin(x) - x * Math.Cos(x)) / (x * x * x);
        double T = BBKS(q);
        return q*q*q * T * T * w * w;
    }
    
    static double P(double k)
    {
        double q = k;
        double T = BBKS(q);
        return A * q * T * T;
    }
            
    static double BBKS(double k)
    {
        double q = k / h / h;
        double a0 = 2.34 * q;
        double a1 = 3.89 * q;
        double a2 = 16.1 * q;
        double a3 = 5.46 * q;
        double a4 = 6.71 * q;
        double T = Math.Log(1.0 + a0) / a0 / Math.Sqrt(Math.Sqrt(1.0 + a1 + a2*a2 + a3*a3*a3 + a4*a4*a4*a4));
            
        return T;
    }
    
    static void Resample(ref double[] sx, ref double[] sy, ref double[] sz, int n)
    {
        int n2 = n / 2;
        int n2tot = n2*n2*n2;
        double[] tx = new double[n2tot];
        double[] ty = new double[n2tot];
        double[] tz = new double[n2tot];
        Console.Error.WriteLine("n = {0}, n2 = {1}, n2tot = {2}", n, n2, n2tot);
        for (int i = 0; i < n2; i++)
        {
            for (int j = 0; j < n2; j++)
            {
                for (int k = 0; k < n2; k++)
                {
                    int si = 2 * i;
                    int sj = 2 * j;
                    int sk = 2 * k;
                    int lefti, leftj, leftk;
                    if (i == 0) lefti = n - 1;
                    else lefti = si - 1;
                    if (j == 0) leftj = n - 1;
                    else leftj = sj - 1;
                    if (k == 0) leftk = n - 1;
                    else leftk = sk - 1;
                    
                    int righti = si + 1;
                    int rightj = sj + 1;
                    int rightk = sk + 1;
                    
                    tx[k + n2 * (j + n2 * i)] = (sx[sk+n*(sj+n*si)] + 
                                                sx[sk+n*(sj+n*lefti)] +
                                                sx[sk+n*(sj+n*righti)] +
                                                sx[sk+n*(leftj+n*si)] +
                                                sx[sk+n*(rightj+n*si)] +
                                                sx[leftk+n*(sj+n*si)] +
                                                sx[rightk+n*(sj+n*si)]) / 7.0;
                                                
                    ty[k + n2 * (j + n2 * i)] = (sy[sk+n*(sj+n*si)] + 
                                                sy[sk+n*(sj+n*lefti)] +
                                                sy[sk+n*(sj+n*righti)] +
                                                sy[sk+n*(leftj+n*si)] +
                                                sy[sk+n*(rightj+n*si)] +
                                                sy[leftk+n*(sj+n*si)] +
                                                sy[rightk+n*(sj+n*si)]) / 7.0;
                                                
                    tz[k + n2 * (j + n2 * i)] = (sz[sk+n*(sj+n*si)] + 
                                                sz[sk+n*(sj+n*lefti)] +
                                                sz[sk+n*(sj+n*righti)] +
                                                sz[sk+n*(leftj+n*si)] +
                                                sz[sk+n*(rightj+n*si)] +
                                                sz[leftk+n*(sj+n*si)] +
                                                sz[rightk+n*(sj+n*si)]) / 7.0;
                                                
                }
            }
        }
        
        sx = tx;
        sy = ty;
        sz = tz;
    }
}

public class RandomGaussian : Random
    {
                   
        private static bool have_next;
        private static double next;
        
        public RandomGaussian() : base()
        {
            next = 0.0;
            have_next = false;
        }
        
        public RandomGaussian(int seed) : base(seed)
        {
            next = 0.0;
            have_next = false;
        }
        
        protected override double Sample()
        {
            if (have_next)
            {
                have_next = false;
                return next;
            }
            else
            {
                double x = 0.0;
                double y = 0.0;
                while (x == 0.0 || y == 0.0)
                {
                    x = base.Sample();
                    y = base.Sample();
                }
                next = Math.Sqrt(-2.0 * Math.Log(x)) * Math.Sin(2.0 * Math.PI * y);
                
                have_next = true;
                return Math.Sqrt(-2.0 * Math.Log(x)) * Math.Cos(2.0 * Math.PI * y);
            }
        }
    }
